from promptAdmin.prompt_app.LLM.AzureOpenai import AzureOpenai
from promptAdmin.prompt_app.LLM.Ollama import OllamaModel
from promptAdmin.prompt_app.utils.GetChat import process_user_input, process_user_input22
from promptAdmin.prompt_app.utils.GetChatChain import get_chat_chain_memory, get_chat_chain_RetrievalQA
from promptAdmin.prompt_app.utils.Vectorstore import extract_text_from_PDF, split_content_into_chunks, \
    read_text_from_txt, split_content_into_chunks_document
from promptAdmin.prompt_app.vector_databases.Chroma import ChromaDatabase
from promptAdmin.prompt_app.vector_databases.azuresearch import AzureSearchDatabase
from promptAdmin.prompt_app.vector_databases.faiss import FaissDatabase


def s_start_prompt(data):
    # 实现复杂的业务逻辑
    return "Result of complex logic : " + data['param1']


def s_store_file(files):
    # 实例化对象（策略模式）
    model = AzureOpenai()
    vector_store = AzureSearchDatabase()

    # files = data.get('files')
    # 4. 获取PDF文档内容（文本）
    # texts = extract_text_from_PDF(files)
    texts = read_text_from_txt(files)
    # 5. 将获取到的文档内容进行切分
    content_chunks = split_content_into_chunks(texts)
    # 6. 对每个chunk计算embedding，并存入到向量数据库
    #     6.1 根据model_type和model_name创建embedding model对象
    embedding_model = model.get_embedding()
    #     6.2 创建向量数据库对象，并将文本embedding后存入到里面
    vector_store = vector_store.save_chunks_into_vectorstore(content_chunks, embedding_model)

    return 'success'


def s_get_chat_chain(data):
    # 实例化对象（策略模式）
    model = AzureOpenai()
    vector_store = AzureSearchDatabase()
    #
    embedding_model = model.get_embedding()
    # 向量数据库
    vector_store = vector_store.get_vectorstore(embedding_model)
    # 相似搜索
    # chat_chain.similarity_search(data['user_input'], vector_store)

    # 获取模型
    llm = model.get_model()
    # 创建对话chain
    chat_chain = get_chat_chain_memory(llm, vector_store)
    # 对话
    ret = process_user_input(chat_chain, data['user_input'])
    return ret



def s_store_file_ollama(files):
    # 实例化对象（策略模式）
    model = OllamaModel()
    vector_store = FaissDatabase()

    # files = data.get('files')
    # 4. 获取PDF文档内容（文本）
    texts = read_text_from_txt(files)
    # 5. 将获取到的文档内容进行切分
    # content_chunks = split_content_into_chunks(texts)
    content_chunks = split_content_into_chunks_document(texts)

    # 6. 对每个chunk计算embedding，并存入到向量数据库
    #     6.1 根据model_type和model_name创建embedding model对象
    embedding_model = model.get_embedding()
    #     6.2 创建向量数据库对象，并将文本embedding后存入到里面
    vector_store = vector_store.save_chunks_into_vectorstore(content_chunks, embedding_model)

    return


def s_get_chat_chain_ollama(data):
    # 实例化对象（策略模式）
    model = OllamaModel()
    vector_store = FaissDatabase()
    # 嵌入
    embedding_model = model.get_embedding()
    # 向量数据库
    vector_store = vector_store.get_vectorstore(embedding_model)
    # 获取模型
    llm = model.get_model()
    # 创建对话chain
    chat_chain = get_chat_chain_RetrievalQA(llm, vector_store)
    # 对话
    ret = process_user_input22(chat_chain, data['user_input'])
    return ret


def s_store_file_ollama_chroma(files):
    # 实例化对象（策略模式）
    model = OllamaModel()
    vector_store = ChromaDatabase()

    # files = data.get('files')
    # 4. 获取PDF文档内容（文本）
    texts = read_text_from_txt(files)
    # 5. 将获取到的文档内容进行切分
    content_chunks = split_content_into_chunks(texts)
    # content_chunks = split_content_into_chunks_document(texts)

    # 6. 对每个chunk计算embedding，并存入到向量数据库
    #     6.1 根据model_type和model_name创建embedding model对象
    embedding_model = model.get_embedding()
    #     6.2 创建向量数据库对象，并将文本embedding后存入到里面
    vector_store = vector_store.save_chunks_into_vectorstore(content_chunks, embedding_model)

    return


def s_get_chat_chain_ollama_chroma(data):
    # 实例化对象（策略模式）
    model = OllamaModel()
    vector_store = ChromaDatabase()
    # 嵌入
    embedding_model = model.get_embedding()
    # 向量数据库
    vector_store = vector_store.get_vectorstore(embedding_model)
    # 获取模型
    llm = model.get_model()
    # 创建对话chain
    chat_chain = get_chat_chain_RetrievalQA(llm, vector_store)
    # 对话
    ret = process_user_input22(chat_chain, data['user_input'])
    return ret

